"""
auther: B24080525 Chenhsi Duan
"""

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

def train_model(train_pixel, train_label, model):
    if model is None:
        model = Sequential([
            Dense(128, activation='relu', input_shape=(784,)),
            Dense(64, activation='relu'),
            Dense(32, activation='relu'),
            Dense(10, activation='softmax')
        ])
        model.compile(optimizer='adam',
            loss='sparse_categorical_crossentropy',
            metrics=['accuracy'
        ])
    model.fit(train_pixel, train_label, epochs=10, batch_size=32, validation_split=0.1)
    model.save('mnist_withoutConv.keras')

train_data = np.loadtxt('train.csv', delimiter=',', skiprows=1)
test_data = np.loadtxt('test.csv', delimiter=',', skiprows=1)
train_label = train_data [:, 0]
train_pixel = train_data [:, 1:] / 255.0
test_pixel = test_data [:, :] / 255.0
try:
    model = tf.keras.models.load_model('mnist_withoutConv.keras')
except:
    train_model(train_pixel, train_label, None)
    model = tf.keras.models.load_model('mnist_withoutConv.keras')
model.summary()
predictions = model.predict(test_pixel)
while 1:
    try:
        i = int(input())
    except:
        print("Invalid value, please try again")
        continue
    if i == -1:
        exit(0)
    elif i == -2:
        train_model(train_pixel, train_label, model)
    elif i < 0:
        print("Invalid value, please try again")
        continue
    else:
        try:
            result = np.argmax(predictions[i])
            plt.imshow(test_pixel[i, :].reshape(28, 28), cmap='gray')
        except:
            print("Index out of range, please try again")
        plt.show()
        print(result)
